# -*- coding: utf-8 -*-
"""
LLM多线程压力测试

测试模型在并发情况下的响应速度和稳定性
"""
import sys
from pathlib import Path
from loguru import logger
import json
import time
import pandas as pd
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

# 添加项目路径
project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

from services import ConversationManager

# 线程锁，用于安全地记录结果
results_lock = threading.Lock()

# 测试消息列表（简化版，用于压力测试）
TEST_MESSAGES = [
    "把k改成0.5",
    "把表层出流系数改成0.6",
    "把下层蓄水容量改成50",
    "把A改成3.5",
    "把初始表层蓄水量改成10",
    "把K改成0.7，把A改成4.0",
    "把S0改成30",
    "把D0改成100",
]


def run_single_test(manager, session_id, message, thread_id, test_id):
    """运行单个测试"""
    start_time = time.time()
    
    try:
        # 发送消息
        response = manager.continue_conversation(
            session_id=session_id,
            user_message=message
        )
        
        elapsed_time = time.time() - start_time
        
        # 检查结果
        success = response and response.get("success", False)
        understood = response.get("understood", False) if response else False
        has_changes = bool(response.get("parameter_changes")) if response else False
        error_msg = response.get("message", "未知错误") if not success else ""
        
        result = {
            "thread_id": thread_id,
            "test_id": test_id,
            "message": message,
            "success": success,
            "understood": understood,
            "has_changes": has_changes,
            "elapsed_time": elapsed_time,
            "error": error_msg if not success else "",
            "timestamp": datetime.now().isoformat()
        }
        
        logger.info(f"[线程{thread_id}] 测试{test_id}: {message[:20]}... - {elapsed_time:.2f}s - {'✅' if success else '❌'}")
        
        return result
        
    except Exception as e:
        elapsed_time = time.time() - start_time
        logger.error(f"[线程{thread_id}] 测试{test_id} 异常: {str(e)}")
        
        return {
            "thread_id": thread_id,
            "test_id": test_id,
            "message": message,
            "success": False,
            "understood": False,
            "has_changes": False,
            "elapsed_time": elapsed_time,
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }


def stress_test(session_id, num_threads=3, tests_per_thread=5):
    """压力测试主函数
    
    Args:
        session_id: 会话ID
        num_threads: 并发线程数
        tests_per_thread: 每个线程的测试次数
    """
    logger.info("=" * 60)
    logger.info("LLM多线程压力测试")
    logger.info("=" * 60)
    logger.info(f"配置:")
    logger.info(f"  并发线程数: {num_threads}")
    logger.info(f"  每线程测试数: {tests_per_thread}")
    logger.info(f"  总测试数: {num_threads * tests_per_thread}")
    logger.info(f"  会话ID: {session_id}")
    logger.info("=" * 60)
    
    # 为每个线程创建独立的ConversationManager
    managers = [ConversationManager() for _ in range(num_threads)]
    
    results = []
    start_time = time.time()
    
    # 创建任务列表
    tasks = []
    for thread_id in range(num_threads):
        for test_id in range(tests_per_thread):
            # 循环使用测试消息
            message = TEST_MESSAGES[test_id % len(TEST_MESSAGES)]
            tasks.append((managers[thread_id], session_id, message, thread_id, test_id))
    
    # 使用线程池执行
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        # 提交所有任务
        futures = {
            executor.submit(run_single_test, *task): task 
            for task in tasks
        }
        
        # 收集结果
        for future in as_completed(futures):
            result = future.result()
            with results_lock:
                results.append(result)
    
    total_time = time.time() - start_time
    
    # 分析结果
    logger.info("\n" + "=" * 60)
    logger.info("测试完成 - 结果分析")
    logger.info("=" * 60)
    
    total_tests = len(results)
    successful_tests = sum(1 for r in results if r["success"])
    understood_tests = sum(1 for r in results if r["understood"])
    
    # 响应时间统计
    times = [r["elapsed_time"] for r in results]
    avg_time = sum(times) / len(times)
    min_time = min(times)
    max_time = max(times)
    
    logger.info(f"\n📊 总体统计:")
    logger.info(f"  总测试数: {total_tests}")
    logger.info(f"  成功数: {successful_tests}")
    logger.info(f"  理解数: {understood_tests}")
    logger.info(f"  成功率: {successful_tests/total_tests*100:.1f}%")
    logger.info(f"  理解率: {understood_tests/total_tests*100:.1f}%")
    
    logger.info(f"\n⏱️  响应时间:")
    logger.info(f"  总耗时: {total_time:.2f}秒")
    logger.info(f"  平均响应: {avg_time:.2f}秒")
    logger.info(f"  最快: {min_time:.2f}秒")
    logger.info(f"  最慢: {max_time:.2f}秒")
    logger.info(f"  吞吐量: {total_tests/total_time:.2f} 请求/秒")
    
    # 按线程统计
    logger.info(f"\n🧵 按线程统计:")
    for thread_id in range(num_threads):
        thread_results = [r for r in results if r["thread_id"] == thread_id]
        thread_success = sum(1 for r in thread_results if r["success"])
        thread_avg_time = sum(r["elapsed_time"] for r in thread_results) / len(thread_results)
        logger.info(f"  线程{thread_id}: {thread_success}/{len(thread_results)} 成功, 平均 {thread_avg_time:.2f}秒")
    
    # 失败的测试
    failed_tests = [r for r in results if not r["success"]]
    if failed_tests:
        logger.warning(f"\n❌ 失败的测试 ({len(failed_tests)}个):")
        for r in failed_tests[:5]:  # 只显示前5个
            logger.warning(f"  [线程{r['thread_id']}] {r['message'][:30]}... - {r['error'][:50]}")
    
    return results, {
        "total_tests": total_tests,
        "successful_tests": successful_tests,
        "understood_tests": understood_tests,
        "total_time": total_time,
        "avg_time": avg_time,
        "min_time": min_time,
        "max_time": max_time,
        "throughput": total_tests / total_time
    }


def save_results_to_csv(results, stats, num_threads, tests_per_thread):
    """保存结果到CSV"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # 详细结果CSV
    detail_csv = Path("storage") / f"stress_test_detail_{timestamp}.csv"
    df_detail = pd.DataFrame(results)
    
    # 重新排列列顺序
    columns_order = [
        "thread_id",
        "test_id",
        "message",
        "success",
        "understood",
        "has_changes",
        "elapsed_time",
        "error",
        "timestamp"
    ]
    df_detail = df_detail[columns_order]
    df_detail.to_csv(detail_csv, index=False, encoding='utf-8-sig')
    
    # 汇总结果CSV
    summary_csv = Path("storage") / f"stress_test_summary_{timestamp}.csv"
    summary_data = {
        "测试配置": [
            f"{num_threads}线程 x {tests_per_thread}测试",
            f"总计 {stats['total_tests']} 个测试"
        ],
        "成功率": [
            f"{stats['successful_tests']}/{stats['total_tests']}",
            f"{stats['successful_tests']/stats['total_tests']*100:.1f}%"
        ],
        "理解率": [
            f"{stats['understood_tests']}/{stats['total_tests']}",
            f"{stats['understood_tests']/stats['total_tests']*100:.1f}%"
        ],
        "总耗时(秒)": [
            f"{stats['total_time']:.2f}",
            ""
        ],
        "平均响应(秒)": [
            f"{stats['avg_time']:.2f}",
            ""
        ],
        "最快响应(秒)": [
            f"{stats['min_time']:.2f}",
            ""
        ],
        "最慢响应(秒)": [
            f"{stats['max_time']:.2f}",
            ""
        ],
        "吞吐量(请求/秒)": [
            f"{stats['throughput']:.2f}",
            ""
        ]
    }
    df_summary = pd.DataFrame(summary_data)
    df_summary.to_csv(summary_csv, index=False, encoding='utf-8-sig')
    
    logger.info(f"\n📊 结果已保存:")
    logger.info(f"  详细结果: {detail_csv}")
    logger.info(f"  汇总结果: {summary_csv}")
    
    return detail_csv, summary_csv


def main():
    """主函数"""
    # 配置参数
    session_id = "de9a990b-59e2-466a-84af-1e0bc1f3c585"
    
    # 可以通过命令行参数调整
    num_threads = int(sys.argv[1]) if len(sys.argv) > 1 else 3
    tests_per_thread = int(sys.argv[2]) if len(sys.argv) > 2 else 5
    
    # 运行压力测试
    results, stats = stress_test(session_id, num_threads, tests_per_thread)
    
    # 保存结果
    detail_csv, summary_csv = save_results_to_csv(results, stats, num_threads, tests_per_thread)
    
    logger.info("\n" + "=" * 60)
    logger.info("压力测试完成！")
    logger.info("=" * 60)
    
    return stats['successful_tests'] == stats['total_tests']


if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        logger.warning("\n\n⚠️  测试被用户中断")
        sys.exit(1)
    except Exception as e:
        logger.error(f"\n\n❌ 测试失败: {str(e)}", exc_info=True)
        sys.exit(1)
